package com.gitee.cashzhang27.test.cloud.oauth.auth.server.extended.component;


import cn.hutool.crypto.SecureUtil;
import cn.hutool.crypto.asymmetric.RSA;
import com.gitee.cashzhang27.test.cloud.oauth.auth.server.extended.comstants.SecurityConstants;
import com.gitee.cashzhang27.test.cloud.oauth.auth.server.extended.util.JwtHelper;
import java.security.KeyPair;
import java.security.PrivateKey;
import java.security.PublicKey;
import java.security.interfaces.RSAPrivateKey;
import java.security.interfaces.RSAPublicKey;
import java.util.Date;
import java.util.HashMap;
import java.util.LinkedHashMap;
import java.util.Map;
import lombok.extern.slf4j.Slf4j;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.data.redis.core.StringRedisTemplate;
import org.springframework.security.jwt.crypto.sign.InvalidSignatureException;
import org.springframework.security.jwt.crypto.sign.MacSigner;
import org.springframework.security.jwt.crypto.sign.RsaSigner;
import org.springframework.security.jwt.crypto.sign.RsaVerifier;
import org.springframework.security.jwt.crypto.sign.SignatureVerifier;
import org.springframework.security.jwt.crypto.sign.Signer;
import org.springframework.security.oauth2.common.DefaultExpiringOAuth2RefreshToken;
import org.springframework.security.oauth2.common.DefaultOAuth2AccessToken;
import org.springframework.security.oauth2.common.DefaultOAuth2RefreshToken;
import org.springframework.security.oauth2.common.ExpiringOAuth2RefreshToken;
import org.springframework.security.oauth2.common.OAuth2AccessToken;
import org.springframework.security.oauth2.common.OAuth2RefreshToken;
import org.springframework.security.oauth2.common.exceptions.InvalidTokenException;
import org.springframework.security.oauth2.common.util.JsonParser;
import org.springframework.security.oauth2.common.util.JsonParserFactory;
import org.springframework.security.oauth2.common.util.RandomValueStringGenerator;
import org.springframework.security.oauth2.provider.OAuth2Authentication;
import org.springframework.security.oauth2.provider.token.AccessTokenConverter;
import org.springframework.security.oauth2.provider.token.DefaultAccessTokenConverter;
import org.springframework.security.oauth2.provider.token.store.JwtClaimsSetVerifier;
import org.springframework.stereotype.Component;
import org.springframework.util.Assert;
import org.springframework.util.Base64Utils;

/**
 * @author Cash Zhang
 * @version v1.0
 * @since 2019/02/10 14:27
 */
@Slf4j
@Component
public class AuthorizationAccessTokenConverter extends
    org.springframework.security.oauth2.provider.token.store.JwtAccessTokenConverter {

  /**
   * Field name for token id.
   */
  public static final String TOKEN_ID = AccessTokenConverter.JTI;

  /**
   * Field name for access token id.
   */
  public static final String ACCESS_TOKEN_ID = AccessTokenConverter.ATI;

  private AccessTokenConverter tokenConverter = new DefaultAccessTokenConverter();

  private JwtClaimsSetVerifier jwtClaimsSetVerifier = new NoOpJwtClaimsSetVerifier();

  private JsonParser objectMapper = JsonParserFactory.create();

  private String verifierKey = new RandomValueStringGenerator().generate();

  private Signer signer = new MacSigner(verifierKey);

  private String signingKey = verifierKey;

  private SignatureVerifier verifier;

  @Autowired
  private StringRedisTemplate redisTemplate;

  /**
   * @param tokenConverter the tokenConverter to set
   */
  @Override
  public void setAccessTokenConverter(AccessTokenConverter tokenConverter) {
    this.tokenConverter = tokenConverter;
  }

  /**
   * @return the tokenConverter in use
   */
  @Override
  public AccessTokenConverter getAccessTokenConverter() {
    return tokenConverter;
  }

  /**
   * @return the {@link JwtClaimsSetVerifier} used to verify the claim(s) in the JWT Claims Set
   */
  @Override
  public JwtClaimsSetVerifier getJwtClaimsSetVerifier() {
    return this.jwtClaimsSetVerifier;
  }

  /**
   * @param jwtClaimsSetVerifier the {@link JwtClaimsSetVerifier} used to verify the claim(s) in the
   * JWT Claims Set
   */
  @Override
  public void setJwtClaimsSetVerifier(JwtClaimsSetVerifier jwtClaimsSetVerifier) {
    Assert.notNull(jwtClaimsSetVerifier, "jwtClaimsSetVerifier cannot be null");
    this.jwtClaimsSetVerifier = jwtClaimsSetVerifier;
  }

  @Override
  public Map<String, ?> convertAccessToken(OAuth2AccessToken token,
      OAuth2Authentication authentication) {
    return tokenConverter.convertAccessToken(token, authentication);
  }

  @Override
  public OAuth2AccessToken extractAccessToken(String value, Map<String, ?> map) {
    return tokenConverter.extractAccessToken(value, map);
  }

  @Override
  public OAuth2Authentication extractAuthentication(Map<String, ?> map) {
    return tokenConverter.extractAuthentication(map);
  }

  /**
   * Unconditionally set the verifier (the verifer key is then ignored).
   *
   * @param verifier the value to use
   */
  @Override
  public void setVerifier(SignatureVerifier verifier) {
    this.verifier = verifier;
  }

  /**
   * Unconditionally set the signer to use (if needed). The signer key is then ignored.
   *
   * @param signer the value to use
   */
  @Override
  public void setSigner(Signer signer) {
    this.signer = signer;
  }

  /**
   * Get the verification key for the token signatures.
   *
   * @return the key used to verify tokens
   */
  @Override
  public Map<String, String> getKey() {
    Map<String, String> result = new LinkedHashMap<>();
    result.put("alg", signer.algorithm());
    result.put("value", verifierKey);
    return result;
  }

  @Override
  public void setKeyPair(KeyPair keyPair) {
    PrivateKey privateKey = keyPair.getPrivate();
    Assert.state(privateKey instanceof RSAPrivateKey, "KeyPair must be an RSA ");
    signer = new RsaSigner((RSAPrivateKey) privateKey);
    RSAPublicKey publicKey = (RSAPublicKey) keyPair.getPublic();
    verifier = new RsaVerifier(publicKey);
    verifierKey =
        "-----BEGIN PUBLIC KEY-----\n" + new String(Base64Utils.encode(publicKey.getEncoded()))
            + "\n-----END PUBLIC KEY-----";
    afterKeyPairSet();
  }

  /**
   * Sets the JWT signing key. It can be either a simple MAC key or an RSA key. RSA keys should be
   * in OpenSSH format, as produced by <tt>ssh-keygen</tt>.
   *
   * @param key the key to be used for signing JWTs.
   */
  @Override
  public void setSigningKey(String key) {
    Assert.hasText(key, "'SigningKey' must not be empty");
    key = key.trim();

    this.signingKey = key;

    if (isPublic(key)) {
      signer = new RsaSigner(key);
      log.info("Configured with RSA signing key");
    } else {
      // Assume it's a MAC key
      this.verifierKey = key;
      signer = new MacSigner(key);
    }
  }

  /**
   * @return true if the key has a public verifier
   */
  private boolean isPublic(String key) {
    return key.startsWith("-----BEGIN");
  }

  /**
   * @return true if the signing key is a public key
   */
  @Override
  public boolean isPublic() {
    return signer instanceof RsaSigner;
  }

  /**
   * The key used for verifying signatures produced by this class. This is not used but is returned
   * from the endpoint to allow resource servers to obtain the key.
   *
   * For an HMAC key it will be the same value as the signing key and does not need to be set. For
   * and RSA key, it should be set to the String representation of the public key, in a standard
   * format (e.g. OpenSSH keys)
   *
   * @param key the signature verification key (typically an RSA public key)
   */
  @Override
  public void setVerifierKey(String key) {
    this.verifierKey = key;
  }

  @Override
  public OAuth2AccessToken enhance(OAuth2AccessToken accessToken,
      OAuth2Authentication authentication) {
    DefaultOAuth2AccessToken result = new DefaultOAuth2AccessToken(accessToken);
    Map<String, Object> info = new LinkedHashMap<>(
        accessToken.getAdditionalInformation());
    String tokenId = result.getValue();
    if (!info.containsKey(TOKEN_ID)) {
      info.put(TOKEN_ID, tokenId);
    } else {
      tokenId = (String) info.get(TOKEN_ID);
    }

    String clientId = authentication.getOAuth2Request().getClientId();
    // 从Redis中取出对应client的公钥私钥对
    String rsaFtfPubKey = redisTemplate.opsForValue()
        .get(SecurityConstants.RSA_KEY_PAIR + clientId + ":" + SecurityConstants.FTF_PUB);
    String rsaFtfPriKey = redisTemplate.opsForValue()
        .get(SecurityConstants.RSA_KEY_PAIR + clientId + ":" + SecurityConstants.FTF_PRI);
    String rsaAppPubKey = redisTemplate.opsForValue()
        .get(SecurityConstants.RSA_KEY_PAIR + clientId + ":" + SecurityConstants.APP_PUB);
    String rsaAppPriKey = redisTemplate.opsForValue()
        .get(SecurityConstants.RSA_KEY_PAIR + clientId + ":" + SecurityConstants.APP_PRI);

    Map<String, String> rsaKeyPair = new HashMap<>();

    rsaKeyPair.put("rsaFtfPubKey", rsaFtfPubKey);
    rsaKeyPair.put("rsaFtfPriKey", rsaFtfPriKey);
    rsaKeyPair.put("rsaAppPubKey", rsaAppPubKey);
    rsaKeyPair.put("rsaAppPriKey", rsaAppPriKey);

    RSA rsa = SecureUtil.rsa(rsaFtfPriKey, rsaFtfPubKey);
    PublicKey publicKey = rsa.getPublicKey();
    PrivateKey privateKey = rsa.getPrivateKey();
    setKeyPair(new KeyPair(publicKey, privateKey));

    result.setAdditionalInformation(info);
    result.setValue(encode(result, authentication, rsaKeyPair));
    OAuth2RefreshToken refreshToken = result.getRefreshToken();
    if (refreshToken != null) {
      DefaultOAuth2AccessToken encodedRefreshToken = new DefaultOAuth2AccessToken(accessToken);
      encodedRefreshToken.setValue(refreshToken.getValue());
      // Refresh tokens do not expire unless explicitly of the right type
      encodedRefreshToken.setExpiration(null);
      try {
        Map<String, Object> claims = objectMapper
            .parseMap(JwtHelper.decode(refreshToken.getValue()).getClaims());
        if (claims.containsKey(TOKEN_ID)) {
          encodedRefreshToken.setValue(claims.get(TOKEN_ID).toString());
        }
      } catch (IllegalArgumentException e) {
      }
      Map<String, Object> refreshTokenInfo = new LinkedHashMap<>(
          accessToken.getAdditionalInformation());
      refreshTokenInfo.put(TOKEN_ID, encodedRefreshToken.getValue());
      refreshTokenInfo.put(ACCESS_TOKEN_ID, tokenId);
      encodedRefreshToken.setAdditionalInformation(refreshTokenInfo);
      DefaultOAuth2RefreshToken token = new DefaultOAuth2RefreshToken(
          encode(encodedRefreshToken, authentication, rsaKeyPair));
      if (refreshToken instanceof ExpiringOAuth2RefreshToken) {
        Date expiration = ((ExpiringOAuth2RefreshToken) refreshToken).getExpiration();
        encodedRefreshToken.setExpiration(expiration);
        token = new DefaultExpiringOAuth2RefreshToken(
            encode(encodedRefreshToken, authentication, rsaKeyPair),
            expiration);
      }
      result.setRefreshToken(token);
    }
    return result;
  }

  @Override
  public boolean isRefreshToken(OAuth2AccessToken token) {
    return token.getAdditionalInformation().containsKey(ACCESS_TOKEN_ID);
  }

  private String encode(OAuth2AccessToken accessToken, OAuth2Authentication authentication,
      Map<String, String> rsaKeyPair) {
    String content;
    try {
      content = objectMapper
          .formatMap(tokenConverter.convertAccessToken(accessToken, authentication));
    } catch (Exception e) {
      throw new IllegalStateException("Cannot convert access token to JSON", e);
    }
    String encryptStr = JwtHelper.encryptStr(content, rsaKeyPair);
    String token = JwtHelper.encode(encryptStr, signer).getEncoded();
    return token;
  }

  public void afterKeyPairSet() {
    if (verifier != null) {
      // Assume signer also set independently if needed
      return;
    }
    SignatureVerifier verifier = new MacSigner(verifierKey);
    try {
      verifier = new RsaVerifier(verifierKey);
    } catch (Exception e) {
      log.warn("Unable to create an RSA verifier from verifierKey (ignoreable if using MAC)");
    }
    // Check the signing and verification keys match
    if (signer instanceof RsaSigner) {
      byte[] test = "com/gitee/cashzhang27/test/boot".getBytes();
      try {
        verifier.verify(test, signer.sign(test));
        log.info("Signing and verification RSA keys match");
      } catch (InvalidSignatureException e) {
        log.error("Signing and verification RSA keys do not match");
      }
    } else if (verifier instanceof MacSigner) {
      // Avoid a race condition where setters are called in the wrong order. Use of
      // == is intentional.
      Assert.state(this.signingKey == this.verifierKey,
          "For MAC signing you do not need to specify the verifier key separately, and if you do it must match the signing key");
    }
    this.verifier = verifier;
  }

  @Override
  public void afterPropertiesSet() {
  }

  private class NoOpJwtClaimsSetVerifier implements JwtClaimsSetVerifier {

    @Override
    public void verify(Map<String, Object> claims) throws InvalidTokenException {
    }
  }

}